mod bencoder;
mod torrent;
mod peer;
mod tracker;

use std::collections::{VecDeque};
use std::net::SocketAddrV4;
use std::path::PathBuf;
use std::sync::{Arc};
use clap::{Parser, Subcommand};
use sha1::{Digest, Sha1};
use tokio::fs::File;
use tokio::io::AsyncWriteExt;
use tokio::sync::{mpsc, Mutex};
use crate::peer::{Peer};

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    #[command(subcommand)]
    command: Command
}

#[derive(Subcommand, Debug)]
#[clap(rename_all = "snake_case")]
enum Command {
    Decode {
        value: String
    },
    Info {
        torrent: PathBuf
    },
    Peers {
        torrent: PathBuf
    },
    Handshake {
        torrent: PathBuf,
        peer: SocketAddrV4
    },
    DownloadPiece {
        #[arg(short)]
        output: PathBuf,
        torrent: PathBuf,
        piece: i32,
    },
    Download {
        #[arg(short)]
        output: PathBuf,
        torrent: PathBuf
    }
}
#[tokio::main]
async fn main() -> anyhow::Result<()>{
    let args = Args::parse();
    match args.command {
        Command::Decode {value} => {
            println!("{}", bencoder::decode_bencoded_value(value.as_str()));
            Ok(())
        },
        Command::Info {torrent} => {
            let torrent = torrent::Torrent::new(torrent);
            println!("Tracker URL: {}", torrent.meta.announce);
            println!("Length: {}", torrent.meta.info.length);
            println!("Info Hash: {}", hex::encode(torrent.info_hash));
            println!("Piece Length: {}", torrent.meta.info.piece_length);
            for hash in torrent.hashes() {
                println!("{}", hash);
            }
            Ok(())
        },
        Command::Peers {torrent} => {
            let torrent = torrent::Torrent::new(torrent);
            let peers = tracker::get_peers(&torrent).await?;
            for peer in peers {
                println!("{}", peer);
            }
            Ok(())
        },
        Command::Handshake {torrent, peer} => {
            let torrent = torrent::Torrent::new(torrent);
            let mut peer = Peer{addr: peer, id: None, session: None };
            peer.handshake(torrent.info_hash).await?;
            Ok(())
        },
        Command::DownloadPiece {torrent, piece, output} => {
            let torrent = torrent::Torrent::new(torrent);
            let peers = tracker::get_peers(&torrent).await?;
            let i_piece = piece;
            for mut peer in peers {
                if peer.handshake(torrent.info_hash.clone()).await.is_err() {
                    continue;
                }
                let pieces = peer.get_pieces().await?;
                if pieces.contains(&piece) {
                    peer.send_interested().await?;
                    peer.get_unchoke().await?;
                    let piece = peer.load_piece(i_piece, &torrent).await?;
                    let hash_from_file = tracker::get_piece_hash(i_piece, &torrent);
                    let mut hasher = Sha1::new();
                    hasher.update(&piece);
                    let real_hash: [u8; 20] = hasher.finalize().try_into().unwrap();
                    anyhow::ensure!(hash_from_file == real_hash);
                    let mut file = File::create(output).await.unwrap();
                    file.write_all(piece.as_slice()).await.unwrap();
                    break;
                }
            }
            Ok(())
        },
        Command::Download {output, torrent} => {
            let torrent = torrent::Torrent::new(torrent);
            let peers = tracker::get_peers(&torrent).await?;
            let pieces = tracker::get_pieces(&torrent);
            let (tx, mut rx) = mpsc::channel(24);
            let pieces: Arc<Mutex<VecDeque<i32>>> = Arc::new(Mutex::new((0..pieces.len() as i32).collect()));
            for mut peer in peers {
                let piece_transmitter = tx.clone();
                let pieces = pieces.clone();
                let torrent = torrent.clone();
                tokio::spawn(async move {
                    if peer.handshake(torrent.info_hash).await.is_err() {
                        println!("Handshake failed");
                        return;
                    }
                    let peer_pieces = peer.get_pieces().await.unwrap();
                    peer.send_interested().await.unwrap();
                    peer.get_unchoke().await.unwrap();
                    loop {
                        if let Some(piece) = pieces.lock().await.pop_front() {
                            println!("Task spawned for {}, piece {piece}", peer.addr);
                            if peer_pieces.contains(&piece) {
                                if let Ok(result) = peer.load_piece(piece, &torrent).await {
                                    println!("Loaded piece {piece} from peer {}", peer.addr);
                                    piece_transmitter.send((piece, result, peer.addr)).await.unwrap();
                                }
                            }
                        }
                        println!("Pieces {:?}", pieces.lock().await);
                        if pieces.lock().await.is_empty() {
                            println!("Piece is empty");
                            return;
                        }
                    }
                });
            }
            drop(tx);
            let mut pieces = vec![];
            while let Some((piece_num, piece, peer)) = rx.recv().await {
                println!("Got piece {piece_num} from {peer}");
                pieces.push((piece_num, piece));
            }
            pieces.sort_by(|el1, el2| el1.0.cmp(&el2.0));
            let mut file = File::create(output.clone()).await.unwrap();
            for piece in pieces {
                file.write_all(piece.1.as_slice()).await.unwrap();
            }
            let size = file.metadata().await.unwrap().len();
            torrent.meta.info.length;
            anyhow::ensure!(torrent.meta.info.length == size as i32);
            Ok(())
        }
    }
}
